import os
import subprocess
import mindspore
from mindspore import context
from thyassist.machine_learning.networks import get_generator
from thyassist.machine_learning.dataloader.load_gan_data import create_dataset
from thyassist.machine_learning.dataloader.load_gan_data import Reporter, save_image, load_ckpt
from launcher import get_project_root


mindspore.set_seed(1)
download_dir = get_project_root()
USE_ORANGE_PI = False
mindspore.set_seed(1)
if os.name == 'nt':
    context.set_context(mode=mindspore.PYNATIVE_MODE, pynative_synchronize=True)
    mindspore.set_device(device_target='CPU')
else:
    try:
        if subprocess.run(['whoami'], capture_output=True, text=True, check=True).stdout.strip() == 'HwHiAiUser':
            context.set_context(mode=context.GRAPH_MODE, jit_config={"jit_level": "O2"})
            mindspore.set_device(device_target='Ascend')
            USE_ORANGE_PI = True
        else:
            context.set_context(mode=context.GRAPH_MODE)
            mindspore.set_device(device_target='Ascend')
    except:
        try:
            context.set_context(mode=context.GRAPH_MODE)
            mindspore.set_device(device_target="GPU")
        except:
            context.set_context(mode=context.GRAPH_MODE)
            mindspore.set_device(device_target="CPU")


def predict(image_size=512):
    """Predict function."""
    rank = 0
    device_num = 1
    G_A = get_generator(image_size=image_size, gl_num=8)
    G_B = get_generator(image_size=image_size, gl_num=8)
    G_A.set_train(True)
    G_B.set_train(True)
    load_ckpt(G_A=G_A, G_B=G_B, is_load_DA=False, is_load_DB=False, is_load_GA=True, is_load_GB=True)
    output_dir = os.path.join(get_project_root(), "figures")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    if not os.path.exists(os.path.join(output_dir, "fake_A")):
        os.makedirs(os.path.join(output_dir, "fake_A"))
    if not os.path.exists(os.path.join(output_dir, "fake_B")):
        os.makedirs(os.path.join(output_dir, "fake_B"))

    # valA
    valA_path = os.path.join(get_project_root(), "gan_datasets", "valA")
    ds = create_dataset(data_path=valA_path, is_train=False, batch_size=1, device_num=1, shuffle=False)
    reporter = Reporter(step_size=ds.get_dataset_size(), is_save_checkpoint=False)
    reporter.start_predict("A to B")
    for data in ds.create_dict_iterator(output_numpy=True):
        img_A = mindspore.Tensor(data["image"])
        path_A = data["image_name"][0]
        path_B = path_A[0:-4] + "_fake_B.jpg"
        fake_B = G_A(img_A)
        save_image(fake_B, os.path.join(output_dir, "fake_B", path_B))
        save_image(img_A, os.path.join(output_dir, "fake_B", path_A))
    reporter.info('save fake_B at %s', os.path.join(output_dir, "fake_B", path_A))
    reporter.end_predict()

    # valB
    valB_path = os.path.join(get_project_root(), "gan_datasets", "valB")
    ds = create_dataset(data_path=valB_path, is_train=False, batch_size=1, device_num=1, shuffle=False)
    reporter.dataset_size = ds.get_dataset_size()
    reporter.start_predict("B to A")
    for data in ds.create_dict_iterator(output_numpy=True):
        img_B = mindspore.Tensor(data["image"])
        path_B = data["image_name"][0]
        path_A = path_B[0:-4] + "_fake_A.jpg"
        fake_A = G_B(img_B)
        save_image(fake_A, os.path.join(output_dir, "fake_A", path_A))
        save_image(img_B, os.path.join(output_dir, "fake_A", path_B))
    reporter.info('save fake_A at %s', os.path.join(output_dir, "fake_A", path_B))
    reporter.end_predict()


def batch_translate(input_dir, output_dir, direction="A2B", image_size=512):
    """将一个目录下的A域转为B域，或将B域转为A域"""
    assert direction in ["A2B", "B2A"], "direction 必须为 'A2B' 或 'B2A'"

    # 初始化生成器
    G_A = get_generator(image_size=image_size, gl_num=8)  # A->B
    G_B = get_generator(image_size=image_size, gl_num=8)  # B->A
    G_A.set_train(False)
    G_B.set_train(False)

    # 加载checkpoint
    load_ckpt(G_A=G_A, G_B=G_B,
              is_load_DA=False, is_load_DB=False,
              is_load_GA=True, is_load_GB=True)

    # 创建数据集
    ds = create_dataset(data_path=input_dir, is_train=False,
                        batch_size=1, device_num=1, shuffle=False)

    # 输出目录准备
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    print(f"============ Start Translation ({direction}) ============")
    for data in ds.create_dict_iterator(output_numpy=True):
        img = mindspore.Tensor(data["image"])
        filename = data["image_name"][0]

        if direction == "A2B":
            fake = G_A(img)
            # 将图像变模糊
            save_path = os.path.join(output_dir, filename.replace(".jpg", "_to_blurry.jpg"))
        else:
            fake = G_B(img)
            # 将图像变清晰
            save_path = os.path.join(output_dir, filename.replace(".jpg", "_to_clear.jpg"))

        save_image(fake, save_path)
        print(f"Saved: {save_path}")
    print(f"============ End Translation ({direction}) ============")


if __name__ == "__main__":
    predict()
